
"""
66服务器上部署的服务：智能取数接口、文件上传接口
"""

import requests
from fake_useragent import UserAgent
from loguru import logger


class ServiceProvider(object):

    def __init__(self):
        self.token = None
        self.refresh_token()

    def refresh_token(self):
        login_url = "http://172.18.0.66/dianda/uis/oauth2/token"
        headers = {
            "Authorization": "Basic dS1jbGllbnQ6MTIzNDU2",
            "User-Agent": UserAgent().random,
        }
        data = {
            "grant_type": "password",
            "username": "superuser",
            "password": "Maiyue!56789@"
        }

        try_times = 5
        while try_times != 0:
            try:
                req = requests.post(login_url, headers=headers, data=data, verify=False)
                result = req.json()['data']
                self.token = f'{result["token_type"]} {result["access_token"]}'
                return
            except Exception as e:
                logger.error(e)
                try_times -= 1

    def get_data(self, query: str):
        """AI 问数查询。"""
        query_url = "http://172.18.0.66/dianda/biw/rpc-api/ai/biwQuery"
        headers = {
            "authorization": self.token,
            "User-Agent": UserAgent().random,
        }
        data = {
            'query': query
        }
        req = requests.post(query_url, headers=headers, data=data, verify=False)
        result = req.json()['data']
        if not isinstance(result, dict):
            return result
        if result.get('msg', '') == '请重新登录系统':
            self.refresh_token()
            return self.get_data(query)
        ask_prompt = result.get('ask_prompt', '')
        if ask_prompt != '':
            return ask_prompt
        if 'aiResult' in result:
            datasets = [
                {key: dataset.get(key, None) for key in ['description', 'rows', 'dimensions']}
                for dataset in result['aiResult']['datasets']
            ]
            datasets = [dataset for dataset in datasets if len(dataset.get('rows')) != 0]
            return datasets
        elif 'dataQueryReq' in result:
            indicator_table = result['dataQueryReq'].get('indicatorTable', None)
            if indicator_table is not None:
                ask_prompt = (f"与此问题相关的表名为: {indicator_table.get('name', '')}, "
                              f"字段描述: {indicator_table.get('dimensionColumnSimpleDisplay', '')}")
                return ask_prompt
        else:
            return result

    def upload_file(self, local_file):
        """文件上传。"""
        upload_url = "http://172.18.0.66/dianda/fs/api/attachment/upload"
        headers = {
            "authorization": self.token
        }
        files = {
            'file': open(local_file, 'rb')
        }
        req = requests.post(url=upload_url,  headers=headers, files=files)
        result = req.json()['data']
        if not isinstance(result, dict):
            return result
        if result.get('msg', '') == '请重新登录系统':
            self.refresh_token()
            return self.upload_file(local_file)
        file_url = f'http://172.18.0.66{result["url"]}'
        return file_url

service_provider = ServiceProvider()